;; Multiarray implementation for nested vectors

;; by Konrad Hinsen
;; last updated April 15, 2010

;; Copyright (c) Konrad Hinsen, 2009-2010. All rights reserved.  The use
;; and distribution terms for this software are covered by the Eclipse
;; Public License 1.0 (http://opensource.org/licenses/eclipse-1.0.php)
;; which can be found in the file epl-v10.html at the root of this
;; distribution.  By using this software in any fashion, you are
;; agreeing to be bound by the terms of this license.  You must not
;; remove this notice, or any other, from this software.

(use 'nstools.ns)
(ns+ multiarray.nested-vectors
  (:clone multiarray.utility)
  (:use [multiarray.protocol]))

(set! *warn-on-reflection* true)

;; Array constructors

(defn zeros
  ([shape]
     (zeros shape Object))
  ([shape element-class]
     (assert (pos? (count shape)))
     (assert (>= (apply min shape) 0))
     (let [[f & r] shape]
       (vec (cond (zero? f) (if (seq r)
			      (throw (Exception. (str "shape " shape " is not "
						      "representable")))
			      [])
		  (nil? r)  (repeat f 0)
		  :else     (repeat f (zeros r element-class)))))))

(defn ma-range
  ([end]
     (ma-range 0 end 1))
  ([start end]
     (ma-range start end 1))
  ([start end step]
     (vec (range start end step))))

(defn- from-elements
  [shape elements]
  (let [rank (count shape)]
    (cond (zero? rank) (first elements)
	  (= rank 1)   (vec elements)
	  :else        (let [[n & r] shape]
			 (if (zero? n)
			   (if (seq r)
			     (throw (Exception. (str "shape " shape " is not "
						     "representable")))
			     [])
			   (let [nrest (quot (count elements) n)]
			     (if (zero? nrest)
			       (vec (repeat n (from-elements r elements)))
			       (vec (map (partial from-elements r)
					 (partition nrest elements))))))))))

(defn into-ma
  [object]
  (from-elements (shape object)
		 (elements object)))

;;
;; Implementation of multiarray.protocol
;;

(defn- shape-as-list
  [x]
  (if (vector? x)
    (let [sub-shapes (map shape-as-list x)]
      (if (all= sub-shapes)
        (cons (count x) (first sub-shapes))
	(throw (Exception. "incorrect vector nesting"))))
    (list)))

(defn nv-reshape
  [x shape-vec]
  (let [n (apply * shape-vec)]
    (from-elements shape-vec
		   (take n (cycle (flatten x))))))

(defn nv-sample 
  [x #^Integer dim #^Integer start #^Integer end #^Integer step]
  (assert (vector? x))
  (assert (not (zero? step)))
  (if (zero? dim)
    (let [n     (count x)
	  back? (neg? step)
	  start (clip-index-to-size n start back?)
	  end   (clip-index-to-size n end back?)]
      (if (pos? step)
	(vec (take-nth step (subvec x start end)))
	(vec (reverse (take-nth (- step) (subvec x (inc end) (inc start)))))))
    (vec (map #(nv-sample % (dec dim) start end step) x))))

(defn nv-slice 
  [x #^Integer dim #^Integer index]
  (assert (vector? x))
  (let [n     (nth (shape x) dim)
	index (check-index-for-size n index)]
    (if (zero? dim)
      (get x index)
      (vec (map #(nv-slice % (dec dim) index) x)))))

(defn- nv-transpose-unchecked
  [x dim-permutation]
  (if  (< (count dim-permutation) 2)
    x
    (let [d         (first dim-permutation)
	  rest-dims (map (fn [i] (if (> i d) (dec i) i))
			 (rest dim-permutation))
	  dth-first (if (zero? d)
		      x
		      (map (partial nv-slice x d) (range (nth (shape x) d))))]
      (vec (map #(nv-transpose-unchecked % rest-dims) dth-first)))))

(defn nv-transpose
  ([x]
   (nv-transpose-unchecked x (range (dec (rank x)) -1 -1)))
  ([x dim-permutation]
   (assert (= (sort dim-permutation) (range (rank x))))
   (nv-transpose-unchecked x dim-permutation)))


(extend clojure.lang.IPersistentVector

  MultiArrayProtocol
  (merge multiarray-method-defaults
	 {:shape 
	    (fn [x] (vec (shape-as-list x)))
	  :elements
	    (fn [x]
	      (if (= (count (shape-as-list x)) 1)
		x
		(vec (flatten x))))
	  :reshape
            nv-reshape
	  :sample
	    nv-sample
	  :slice
	    nv-slice
	  :transpose
	    nv-transpose})

  MultiArrayLowLevelProtocol
  (merge multiarray-lowlevel-method-defaults
	 {:element-class
  	    (comp common-element-class flatten)}))


(set! *warn-on-reflection* false)
